import subprocess as sp
import numpy as np
import json
import sys
import os

def eval_cmp(src_path, dst_path, m, n, k):
    my = np.fromfile(src_path, dtype=np.int32).reshape((m, n))
    gt = np.fromfile(dst_path, dtype=np.int32).reshape((m, n))
    diff = np.abs(my-gt)
    equal = np.all(gt==my)
    if not equal:
        print(f"## m {m} n {n} k {k} not equal {os.path.basename(src_path)} max_abs_diff {np.argmax(diff)}: {np.max(diff)}")
        np.set_printoptions(threshold=np.inf)
        print(f"{os.path.basename(src_path)} \n {my}")
        print(f"{os.path.basename(dst_path)} \n {gt}")
        sys.exit(-1)
    else:
        print(f"## m {m} n {n} k {k} equal")

cast_path = os.path.realpath(sys.argv[1])
cmd="compute-sanitizer --tool memcheck ./mp_gemm"
# cmd="./mp_gemm"

gt_prefix = "../../eval_gt/"
f = open(cast_path,"r") 
lines = f.readlines()
for line in lines:
    its = json.loads(line)
    m, n, k = its["m"], its["n"], its["k"]
    A_gt_path, B_gt_path, C_gt_path = its["A_path"], its["B_path"], its["C_path"]
    test_cmd = cmd.split()+[f"{m}", f"{n}", f"{k}", "1"]
    C_my_path = gt_prefix+"./mp_gemm_w1a8.bin"
    sp.run(test_cmd + [gt_prefix+A_gt_path, gt_prefix+B_gt_path, C_my_path])
    eval_cmp(C_my_path, gt_prefix+C_gt_path, m, n, k)

